use clap::Parser;
use ordered_float::NotNan;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
use rimec::{
    imec::{
        common::{DataOfInterest, NextTokenDistributionProvider, compute_marginals},
        decode::decode,
        encode::encode,
    },
    mec::DiscreteDistribution,
};
use tokio::{sync::Mutex, task};

use std::sync::Arc;
use std::{cell::RefCell, fs::File, io};

use itertools::sorted;

#[derive(clap::Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    #[arg(short, long)]
    output: String,
}

extern crate slog;
extern crate slog_async;
extern crate slog_stream;
extern crate slog_term;

use slog::Drain;
use slog::o;

pub struct AsciidocReport<W: io::Write> {
    io: RefCell<W>,
}

pub enum ReportElement {
    Table {
        title: Option<String>,
        headings: Vec<String>,
        contents: Vec<Vec<String>>,
    },
    Paragraph(String),
    Heading(u8, String),
}

impl<W> AsciidocReport<W>
where
    W: io::Write,
{
    pub fn new(io: W) -> AsciidocReport<W> {
        AsciidocReport {
            io: RefCell::new(io),
        }
    }

    pub fn write(&mut self, e: ReportElement) {
        match e {
            ReportElement::Table {
                title,
                headings,
                contents,
            } => {
                if let Some(title) = title {
                    let _ = writeln!(self.io.borrow_mut(), ".{}", title);
                };

                let _ = writeln!(
                    self.io.borrow_mut(),
                    "[cols=\"{}\",stripes=even]",
                    (0..headings.len())
                        .map(|_i| "1")
                        .collect::<Vec<_>>()
                        .join(",")
                );
                let _ = writeln!(self.io.borrow_mut(), "|===");

                headings.iter().for_each(|cell| {
                    let _ = writeln!(self.io.borrow_mut(), "|{}", cell);
                });
                let _ = writeln!(self.io.borrow_mut());

                contents.iter().for_each(|row| {
                    row.iter().for_each(|cell| {
                        let _ = writeln!(self.io.borrow_mut(), "|{}", cell);
                    });
                    let _ = writeln!(self.io.borrow_mut());
                });

                let _ = writeln!(self.io.borrow_mut(), "|===");
                let _ = writeln!(self.io.borrow_mut());
            }
            ReportElement::Paragraph(string) => {
                let _ = writeln!(self.io.borrow_mut(), "{}\n", string);
            }
            ReportElement::Heading(u8, string) => {
                let prefix = "#".repeat(u8 as usize);
                let _ = writeln!(self.io.borrow_mut(), "{} {}\n", prefix, string);
            }
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let args = Args::parse();

    let alphabet: Vec<&str> = vec!["A", "B", "C"];

    let mut provider: DiscreteDistribution<&str> = DiscreteDistribution {
        d: vec![0.3, 0.3, 0.4],
        labels: Some(alphabet.clone()),
    };

    let input_space_labels: Vec<usize> = (0..=7).collect();

    let input: Vec<usize> = vec![1, 0, 1, 2, 3, 3, 1, 0, 7, 5];

    let output_doc_file = File::create(args.output).expect("Couldn't open output file");
    let output_doc = Arc::new(Mutex::new(AsciidocReport::new(output_doc_file)));

    let decorator = slog_term::TermDecorator::new().build();
    let drain = slog_term::FullFormat::new(decorator).build().fuse();
    let drain = slog_async::Async::new(drain)
        .overflow_strategy(slog_async::OverflowStrategy::Block)
        .build()
        .fuse();

    let mut logger = slog::Logger::root(drain, o!());

    write_step_by_step_document(
        output_doc,
        // alphabet,
        &mut provider,
        &input_space_labels,
        input,
        &mut logger,
    )
    .await
}

// enum Mode {
//     Start,
// }

macro_rules! concat ( ($part1: expr, $part2: expr) => {
    $part1.into_iter().chain($part2.into_iter()).collect::<Vec<_>>()
});

async fn write_step_by_step_document<'b, W: io::Write + 'b + 'static + Send>(
    output_doc: Arc<Mutex<AsciidocReport<W>>>,
    // alphabet: Vec<&'b str>,
    provider: &'b mut impl NextTokenDistributionProvider<&'b str, ()>,
    input_space_labels: &[usize],
    input: Vec<usize>,
    logger: &'b mut slog::Logger,
) -> Result<(), anyhow::Error> {
    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<DataOfInterest>();

    // let mut mode = Mode::Start;

    // let alphabet_for_logging = alphabet.iter().map(|x| x.to_string()).collect::<Vec<_>>();

    let input_symbols = input_space_labels
        .iter()
        .map(|x| x.to_string())
        .collect::<Vec<_>>();

    let logging = task::spawn(async move {
        while let Some(data_element) = rx.recv().await {
            let elements = match data_element {
                DataOfInterest::DecodeStart() => {
                    vec![ReportElement::Heading(2, "Decode Start".to_string())]
                }
                DataOfInterest::EncodeStart() => {
                    vec![ReportElement::Heading(2, "Encode Start".to_string())]
                }
                DataOfInterest::IStar(istar) => {
                    vec![ReportElement::Paragraph(format!("**istar = {}**", istar))]
                }
                DataOfInterest::EncodeIteration(n) => {
                    vec![ReportElement::Heading(3, format!("Encode iteration {}", n))]
                }
                DataOfInterest::FinalResult(string, input_string) => {
                    vec![
                        ReportElement::Paragraph(format!("Final result {}", string)),
                        ReportElement::Paragraph(format!("Input was {}", input_string)),
                    ]
                }
                DataOfInterest::EncodeResult(string) => {
                    vec![ReportElement::Paragraph(format!(
                        "Encoded string: {}",
                        string
                    ))]
                }
                DataOfInterest::GammaTable(gammaj) => {
                    let (rep, req) = compute_marginals(&gammaj);
                    let firstzero_p = if rep
                        .is_sorted_by(|x, y| y.partial_cmp(x) == Some(std::cmp::Ordering::Less))
                    {
                        req.iter()
                            .enumerate()
                            .find(|(_i, x)| (**x).abs() <= 1.0 * f64::EPSILON)
                            .map(|x| x.0)
                            .unwrap_or(rep.len())
                    } else {
                        req.len()
                    };
                    vec![ReportElement::Table {
                        title: Some("gamma table".to_string()),
                        headings: concat!(
                            vec!["".to_string()],
                            req.iter()
                                .take(firstzero_p)
                                .map(|x| format!("{:.4}", x))
                                .collect::<Vec<_>>()
                        ),
                        contents: (0..rep.len())
                            .map(|i| {
                                concat!(
                                    vec![format!("{:.4}", rep[i])],
                                    (0..std::cmp::min(firstzero_p, req.len()))
                                        .map(|j| {
                                            format!(
                                                "{:.4}",
                                                gammaj
                                                    .iter()
                                                    .find(|(_r, (i_, j_))| (*i_, *j_) == (i, j))
                                                    .map_or_else(|| 0.0, |x| x.0)
                                            )
                                        })
                                        .collect::<Vec<_>>()
                                )
                            })
                            .collect::<Vec<_>>(),
                    }]
                }

                DataOfInterest::MuTable(rows, maybe_istar) => {
                    vec![ReportElement::Table {
                        title: Some("mu table".to_string()),
                        headings: concat!(input_symbols.clone(), vec!["entropy".to_string()]),
                        contents: rows
                            .iter()
                            .enumerate()
                            .map(|(i, (labels, row, entropy))| {
                                let floats = sorted(row.iter().enumerate().map(|(ix, x)| {
                                    (
                                        labels[ix].clone(),
                                        std::cmp::Reverse(NotNan::new(*x).unwrap()),
                                    )
                                }))
                                .map(|(_, x)| x)
                                .map(|x| {
                                    if Some(i) == maybe_istar {
                                        format!("**{:.4}**", x.0)
                                    } else {
                                        format!("{:.4}", x.0)
                                    }
                                });

                                let entropy_field = vec![format!("{}", entropy)].into_iter();

                                floats.into_iter().chain(entropy_field).collect::<Vec<_>>()
                            })
                            .collect::<Vec<_>>(),
                    }]
                }
            };

            for element in elements {
                output_doc.lock().await.write(element);
            }
        }
    });

    let mut rng = ChaCha8Rng::seed_from_u64(4223);

    let output = encode(
        &mut rng,
        &input[..],
        input_space_labels,
        provider,
        (),
        Some(logger),
        Some(&tx),
    )
    .await?;

    let regenerated_input = decode(
        &mut output.into_iter(),
        input_space_labels,
        provider,
        input.len(),
        (),
        Some(logger),
        Some(&tx),
    )
    .await?;

    tx.send(DataOfInterest::FinalResult(
        format!("{:?}", regenerated_input),
        format!("{:?}", input),
    ))
    .unwrap();

    drop(tx);

    let _ = tokio::join!(logging);

    Ok(())
}
